# set up environment ----
library(data.table)
library(nnet)
library(RColorBrewer)
library(ggplot2)
library(cowplot)
load("data/analysis_data.RData")
source("R/00-functions.R")
set.seed(1991975877)


respondent_data[n_attention_fails <= 1 & trump_least_preferred == 1 &
    support_2016 != "Trump",
  .N / 833, type3]

# average treatment effect tests ----
# main dataset
respondent_data[n_attention_fails <= 1 & trump_least_preferred == 1 &
    support_2016 != "Trump",
  chisq.test(electability_last, type3)]
respondent_data[n_attention_fails <= 1 & trump_least_preferred == 1 &
    support_2016 != "Trump",
  fisher.test(electability_last, type3)]

# full dataset
respondent_data[, chisq.test(electability_last, type3)]
respondent_data[, fisher.test(electability_last, type3)]

# correlation matrix ----
respondent_data[, cor(cbind(need_for_cognition, campaign_knowledge,
  follows_politics, political_interest))]

# ATEs ----
load("results/predicted_probs_main.RData")
calc_ates <- function(x1, x2) {
  predicted_probs_main[, .(
    ate_p = mean(preferences[eval(parse(text = x1))]) -
      mean(preferences[eval(parse(text = x2))]),
    ate_e = mean(electability[eval(parse(text = x1))]) -
      mean(electability[eval(parse(text = x2))]),
    ate_eu = mean(expected_utility[eval(parse(text = x1))]) -
      mean(expected_utility[eval(parse(text = x2))])
  ), b][, .(
    outcome =  c("preferences", "electability", "expected utility"),
    est = c(mean(ate_p), mean(ate_e), mean(ate_eu)),
    q025 = c(quantile(ate_p, .025), quantile(ate_e, .025),
      quantile(ate_eu, .025)),
    q975 = c(quantile(ate_p, .975), quantile(ate_e, .975),
      quantile(ate_eu, .975)),
    pval = c(
      2 * min(mean(ate_p > 0), mean(ate_p < 0)),
      2 * min(mean(ate_e > 0), mean(ate_e < 0)),
      2 * min(mean(ate_eu > 0), mean(ate_eu < 0))))]
}
calc_ates('order %like% "^E"', 'order %like% "^P"')

# euclidean distance on ternary plot ----
load(file = "results/predicted_probs_main_1.RData")
load(file = "results/predicted_probs_main_0.RData")
X <- merge(predicted_probs_main_1, predicted_probs_main_0,
  by = c("i", "b"))
X[, `:=`(
  bary1d.x = 1/2 * electability.x + expected_utility.x,
  bary2d.x = sqrt(3) / 2 * electability.x,
  bary1d.y = 1/2 * electability.y + expected_utility.y,
  bary2d.y = sqrt(3) / 2 * electability.y)]
X <- X[, sqrt(
  (bary1d.x - bary1d.y) ^  2 +
  (bary2d.x - bary2d.y) ^  2),  .(i, b)]
X[, mean(V1), b][, .(mean(V1), quantile(V1, .025), quantile(V1, .975))]

# marginal effects ----
calc_ates(
  'need_for_cognition == "High"',
  'need_for_cognition == "Low"')
calc_ates(
  'campaign_knowledge == "High"',
  'campaign_knowledge == "Low"')
calc_ates(
  'political_interest == 1',
  'political_interest == 0')
calc_ates(
  'follows_politics == 1',
  'follows_politics == 0')

# candidates ----
load("results/boots_coefs_main.RData")
boot_coefs_main[[1]][, .(
  est = mean(`electability_last:white_male_candidate`),
  q025 = quantile(`electability_last:white_male_candidate`, .025),
  q975 = quantile(`electability_last:white_male_candidate`, .975),
  pval = 2 * min(mean(`electability_last:white_male_candidate` > 0),
    mean(`electability_last:white_male_candidate` < 0)))]
boot_coefs_main[[2]][, .(
  est = mean(`electability_last:straight_white_male_candidate`),
  q025 = quantile(`electability_last:straight_white_male_candidate`, .025),
  q975 = quantile(`electability_last:straight_white_male_candidate`, .975),
  pval = 2 * min(mean(`electability_last:straight_white_male_candidate` > 0),
    mean(`electability_last:straight_white_male_candidate` < 0)))]

